package org.hamster.common;

import java.math.BigInteger;
import java.util.regex.Pattern;

/**
 * A class for handling IPv6 addresses and IPv6 masks.
 * Provides functionality to validate IPv6 addresses and check if an IPv6 is within a given range.
 */
public class IPv6Mask implements IPMaskable {
    
    private final BigInteger startIp;
    private final BigInteger endIp;
    private final int prefixLength;
    
    /**
     * Constructs an IPv6Mask from a string that can be either an IPv6 address or an IPv6 range in CIDR notation.
     * 
     * @param ipString the IPv6 address or IPv6 range in CIDR notation (e.g. "2001:db8::1" or "2001:db8::/32")
     * @throws IllegalArgumentException if the input string is not a valid IPv6 address or CIDR notation
     */
    public IPv6Mask(String ipString) {
        if (ipString == null || ipString.isEmpty()) {
            throw new IllegalArgumentException("IPv6 string cannot be null or empty");
        }
        
        if (ipString.contains("/")) {
            // CIDR notation
            String[] parts = ipString.split("/");
            // Check if we have exactly two parts (IP and prefix)
            if (parts.length != 2) {
                throw new IllegalArgumentException("Invalid CIDR notation: " + ipString);
            }
            
            String ipPart = parts[0];
            int prefixLength;
            try {
                prefixLength = Integer.parseInt(parts[1]);
            } catch (NumberFormatException e) {
                throw new IllegalArgumentException("Invalid prefix length in CIDR notation: " + parts[1]);
            }
            
            if (!IPValidationUtil.isValidIPv6(ipPart)) {
                throw new IllegalArgumentException("Invalid IPv6 address in CIDR notation: " + ipPart);
            }
            
            if (prefixLength < 0 || prefixLength > 128) {
                throw new IllegalArgumentException("Prefix length must be between 0 and 128, got: " + prefixLength);
            }
            
            BigInteger ip = ipv6ToBigInteger(ipPart);
            if (prefixLength == 128) {
                // Special case for /128 - single IP address
                startIp = ip;
                endIp = ip;
            } else {
                BigInteger mask = BigInteger.ONE.shiftLeft(128 - prefixLength).subtract(BigInteger.ONE).shiftLeft(prefixLength);
                startIp = ip.and(mask);
                endIp = startIp.or(BigInteger.ONE.shiftLeft(128 - prefixLength).subtract(BigInteger.ONE));
            }
            this.prefixLength = prefixLength;
        } else {
            // Single IPv6 address
            if (!IPValidationUtil.isValidIPv6(ipString)) {
                throw new IllegalArgumentException("Invalid IPv6 address: " + ipString);
            }
            
            startIp = ipv6ToBigInteger(ipString);
            endIp = startIp;
            this.prefixLength = 128; // For a single IP, the prefix length is 128
        }
    }
    
    /**
     * Converts an IPv6 address string to a BigInteger value.
     * 
     * @param ipv6 the IPv6 address string
     * @return the BigInteger representation of the IPv6 address
     */
    private static BigInteger ipv6ToBigInteger(String ipv6) {
        // Expand compressed format
        String expanded = expandIPv6(ipv6);
        String[] parts = expanded.split(":");
        BigInteger result = BigInteger.ZERO;
        for (int i = 0; i < 8; i++) {
            BigInteger part = new BigInteger(parts[i], 16);
            result = result.shiftLeft(16).add(part);
        }
        return result;
    }
    
    /**
     * Expands a compressed IPv6 address to its full form.
     * 
     * @param ipv6 the IPv6 address string (possibly compressed)
     * @return the expanded IPv6 address string
     */
    private static String expandIPv6(String ipv6) {
        if (ipv6.contains("::")) {
            String[] parts = ipv6.split("::", -1);
            String leftPart = parts[0];
            String rightPart = parts[1];
            
            int leftCount = leftPart.isEmpty() ? 0 : (leftPart.contains(":") ? leftPart.split(":").length : (leftPart.isEmpty() ? 0 : 1));
            int rightCount = rightPart.isEmpty() ? 0 : (rightPart.contains(":") ? rightPart.split(":").length : (rightPart.isEmpty() ? 0 : 1));
            
            int missingCount = 8 - leftCount - rightCount;
            
            StringBuilder sb = new StringBuilder();
            if (!leftPart.isEmpty()) {
                sb.append(leftPart);
                if (!leftPart.endsWith(":")) {
                    sb.append(":");
                }
            }
            
            for (int i = 0; i < missingCount; i++) {
                sb.append("0000");
                if (i < missingCount - 1) {
                    sb.append(":");
                }
            }
            
            if (!rightPart.isEmpty()) {
                if (!rightPart.startsWith(":") && sb.length() > 0) {
                    sb.append(":");
                }
                sb.append(rightPart);
            }
            
            // Handle special case of just "::"
            if (leftPart.isEmpty() && rightPart.isEmpty()) {
                sb.setLength(0);
                for (int i = 0; i < 8; i++) {
                    sb.append("0000");
                    if (i < 7) {
                        sb.append(":");
                    }
                }
            }
            
            ipv6 = sb.toString();
        }
        
        // Pad each part to 4 characters
        String[] parts = ipv6.split(":");
        StringBuilder result = new StringBuilder();
        for (int i = 0; i < parts.length; i++) {
            if (i > 0) {
                result.append(":");
            }
            String part = parts[i];
            while (part.length() < 4) {
                part = "0" + part;
            }
            result.append(part);
        }
        
        // Fill in missing parts if any (when there are less than 8 parts)
        int colons = 0;
        for (char c : ipv6.toCharArray()) {
            if (c == ':') {
                colons++;
            }
        }
        
        while (colons < 7 && parts.length < 8) {
            result.append(":0000");
            colons++;
        }
        
        return result.toString();
    }
    
    /**
     * Converts a BigInteger value to an IPv6 address string.
     * 
     * @param ip the BigInteger representation of the IPv6 address
     * @return the IPv6 address in standard notation
     */
    private static String bigIntegerToIPv6(BigInteger ip) {
        StringBuilder result = new StringBuilder();
        BigInteger mask = new BigInteger("FFFF", 16);
        
        for (int i = 0; i < 8; i++) {
            if (i > 0) {
                result.insert(0, ":");
            }
            BigInteger part = ip.and(mask);
            result.insert(0, String.format("%x", part));
            ip = ip.shiftRight(16);
        }
        
        return compressIPv6(result.toString());
    }
    
    /**
     * Compresses an IPv6 address by replacing the longest sequence of zero groups with ::
     * 
     * @param ipv6 the full IPv6 address string
     * @return the compressed IPv6 address string
     */
    private static String compressIPv6(String ipv6) {
        // Find the longest sequence of consecutive zero groups
        String[] parts = ipv6.split(":");
        int maxStart = -1;
        int maxLen = 0;
        int currentStart = -1;
        int currentLen = 0;
        
        for (int i = 0; i < parts.length; i++) {
            if ("0".equals(parts[i]) || "0000".equals(parts[i])) {
                if (currentStart == -1) {
                    currentStart = i;
                    currentLen = 1;
                } else {
                    currentLen++;
                }
            } else {
                if (currentLen > maxLen && currentLen > 1) {
                    maxStart = currentStart;
                    maxLen = currentLen;
                }
                currentStart = -1;
                currentLen = 0;
            }
        }
        
        // Check at the end
        if (currentLen > maxLen && currentLen > 1) {
            maxStart = currentStart;
            maxLen = currentLen;
        }
        
        // Special case for all zeros
        boolean allZeros = true;
        for (String part : parts) {
            if (!"0".equals(part) && !"0000".equals(part)) {
                allZeros = false;
                break;
            }
        }
        if (allZeros) {
            return "::";
        }
        
        // If we found a sequence to compress
        if (maxStart != -1) {
            StringBuilder result = new StringBuilder();
            for (int i = 0; i < parts.length; i++) {
                if (i == maxStart) {
                    if (i == 0 || result.length() == 0) {
                        result.append(":");
                    }
                    result.append(":");
                    i += maxLen - 1;
                    if (i == parts.length - 1) {
                        result.append(":");
                    }
                } else {
                    if (result.length() > 0) {
                        result.append(":");
                    }
                    // Remove leading zeros but keep at least one digit
                    String part = parts[i].replaceFirst("^0+(?!$)", "");
                    if (part.isEmpty()) {
                        part = "0";
                    }
                    result.append(part);
                }
            }
            return result.toString();
        }
        
        // No compression needed, just remove leading zeros
        StringBuilder result = new StringBuilder();
        for (int i = 0; i < parts.length; i++) {
            if (i > 0) {
                result.append(":");
            }
            // Remove leading zeros but keep at least one digit
            String part = parts[i].replaceFirst("^0+(?!$)", "");
            if (part.isEmpty()) {
                part = "0";
            }
            result.append(part);
        }
        return result.toString();
    }
    
    /**
     * Checks if the given IPv6 string is a valid IPv6 address.
     * 
     * @param ipv6 the IPv6 address string to validate
     * @return true if the IPv6 is valid, false otherwise
     */
    public static boolean isValidIPv6(String ipv6) {
        return IPValidationUtil.isValidIPv6(ipv6);
    }
    
    /**
     * Checks if the given IPv6 address is within the range represented by this IPv6Mask.
     * 
     * @param ipv6 the IPv6 address to check
     * @return true if the IPv6 is valid and within range, false otherwise
     */
    @Override
    public boolean inRange(String ipv6) {
        if (ipv6 == null || !IPValidationUtil.isValidIPv6(ipv6)) {
            return false;
        }
        
        BigInteger ip = ipv6ToBigInteger(ipv6);
        return ip.compareTo(startIp) >= 0 && ip.compareTo(endIp) <= 0;
    }
    
    /**
     * Returns the start IPv6 of the range as a string.
     * 
     * @return the start IPv6 address
     */
    @Override
    public String getStartIp() {
        return bigIntegerToIPv6(startIp);
    }
    
    /**
     * Returns the end IPv6 of the range as a string.
     * 
     * @return the end IPv6 address
     */
    @Override
    public String getEndIp() {
        return bigIntegerToIPv6(endIp);
    }
    
    @Override
    public String toString() {
        if (prefixLength == 128) {
            return bigIntegerToIPv6(startIp);
        } else {
            return bigIntegerToIPv6(startIp) + "/" + prefixLength;
        }
    }
}